#!/bin/bash
###############################################################################
#                                                                             #
#                        Pi-Star Partition Expander                           #
#                                                                             #
#    Version 1.1, Code, Design and Development by Andy Taylor (MW0MWZ).       #
#                                                                             #
#    This script will expand your root file system to fill the SD Card        #
#                                                                             #
###############################################################################

if [ "$(id -u)" != "0" ]; then
  echo -e "You need to be root to run this command...\n"
  exit 1
fi

# Remount read-write
if [ -d /boot/firmware ]; then
  mount -o remount,rw /boot/firmware
else
  mount -o remount,rw /boot
fi
mount -o remount,rw /

# Figure out if this is a systemd system
if command -v systemctl > /dev/null && systemctl | grep -q '\-\.mount'; then
	SYSTEMD=1
elif [ -f /etc/init.d/cron ] && [ ! -h /etc/init.d/cron ]; then
	SYSTEMD=0
else
	echo "Error - Unrecognised init system"
	return 1
fi

# Expand the root partition
if [ $SYSTEMD -eq 1 ]; then
	ROOT_PART=$(mount | sed -n 's|^/dev/\(.*\) on / .*|\1|p')
else
	if ! [ -h /dev/root ]; then
		echo "Error - dont know how to expand this filesystem"
		return 0
	fi
	ROOT_PART=$(readlink /dev/root)
fi

if [ "${ROOT_PART:0:6}" != "mmcblk" ] || [ "${ROOT_PART:8:1}" != "2" ]; then
	echo "Your partition layout is not currently supported by this tool, or your root filesystem is already expanded."
	return 0
fi

LAST_PART_NUM=$(parted /dev/${ROOT_PART:0:7} -ms unit s p | tail -n 1 | cut -f 1 -d:)
if [ $LAST_PART_NUM -ne ${ROOT_PART:8:1} ]; then
	echo "$ROOT_PART is not the last partition. Don't know how to expand"
	return 0
fi
# Get the starting offset of the root partition
PART_START=$(parted /dev/${ROOT_PART:0:7} -ms unit s p | grep "^${ROOT_PART:8:1}" | cut -f 2 -d: | sed 's/[^0-9]//g')
[ "$PART_START" ] || return 1
# Return value will likely be error for fdisk as it fails to reload the
# partition table because the root fs is mounted
fdisk /dev/${ROOT_PART:0:7} <<EOF
p
d
${ROOT_PART:8:1}
n
p
${ROOT_PART:8:1}
$PART_START

p
w
EOF

# now set up an init.d script
cat <<EOF > /etc/init.d/resize2fs_once &&
#!/bin/sh
### BEGIN INIT INFO
# Provides:          resize2fs_once
# Required-Start:
# Required-Stop:
# Default-Start: 3
# Default-Stop:
# Short-Description: Resize the root filesystem to fill partition
# Description:
### END INIT INFO

. /lib/lsb/init-functions

case "\$1" in
  start)
    log_daemon_msg "Starting resize2fs_once" &&
    if [ -d /boot/firmware ]; then
      mount -o remount,rw /boot/firmware
    else
      mount -o remount,rw /boot
    fi
    mount -o remount,rw / &&
    resize2fs /dev/$ROOT_PART &&
    update-rc.d resize2fs_once remove &&
    rm /etc/init.d/resize2fs_once &&
    log_end_msg \$?
    ;;
  *)
    echo "Usage: \$0 start" >&2
    exit 3
    ;;
esac
EOF
chmod +x /etc/init.d/resize2fs_once &&
update-rc.d resize2fs_once defaults &&

# Remount read-only
if [ -d /boot/firmware ]; then
  mount -o remount,ro /boot/firmware
else
  mount -o remount,ro /boot
fi
mount -o remount,ro /

# Finished, so tell the user
echo "Partition expansion has been completed, please reboot to expand the filesystem to fill it"

exit 0
